use std::{collections::HashMap, net::SocketAddr, sync::Arc};

use anyhow::Result;
use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        ConnectInfo, Query,
    },
    response::Response,
    routing::get,
    Extension, Router, Server,
};
use futures_util::{stream::SplitSink, SinkExt, StreamExt, TryStreamExt};
use serde::Deserialize;
use tokio::{net::TcpStream, sync::Mutex};
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use tracing::info;
use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
use urlencoding::encode;

mod config;
mod error;
mod helpers;

use config::Settings;
use error::AppError;
use helpers::*;

type SharedChannels = Mutex<HashMap<SocketAddr, SplitSink<WebSocket, Message>>>;

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::registry()
        .with(fmt::layer())
        .with(EnvFilter::from_default_env())
        .init();

    let settings = Settings::load()?;
    let listening_addr = settings.app_listening_addr()?;
    info!("Listening on: {listening_addr}");
    let channels: SharedChannels = Mutex::new(HashMap::new());

    let app = Router::new()
        .route("/ws/student", get(ws_student))
        .route("/ws/teacher", get(ws_teacher))
        .layer(Extension(Arc::new(channels)))
        .layer(Extension(Arc::new(settings)));

    Server::bind(&listening_addr)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await
        .map_err(Into::into)
}

#[derive(Debug, Deserialize)]
struct Params {
    from: String,
    to: String,
    provider: String,
}

async fn ws_teacher(
    ws: WebSocketUpgrade,
    Query(params): Query<Params>,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    Extension(channels): Extension<Arc<SharedChannels>>,
    Extension(settings): Extension<Arc<Settings>>,
) -> Result<Response, AppError> {
    let provider_settings = settings
        .provider_settings(&params.provider)
        .ok_or(AppError::UnsupportedProvider(params.provider))?;

    let mut provider_url = provider_settings.ws_url()?;
    provider_url = provider_url.join("/ws/voice")?;
    provider_url
        .query_pairs_mut()
        .append_pair("from", encode(&params.from).as_ref());
    provider_url
        .query_pairs_mut()
        .append_pair("to", encode(&params.to).as_ref());

    let (pws, ..) = connect_async(&provider_url)
        .await
        .map_err(|err| AppError::CannotConnectProvider(provider_url, err))?;

    Ok(ws.on_upgrade(move |ws| ws_teacher_handler(addr, ws, pws, channels)))
}

async fn ws_student(
    ws: WebSocketUpgrade,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    Extension(channels): Extension<Arc<SharedChannels>>,
) -> Response {
    ws.on_upgrade(move |ws| ws_student_handler(addr, ws, channels))
}

async fn ws_teacher_handler(
    id: SocketAddr,
    ws: WebSocket,
    pws: WebSocketStream<MaybeTlsStream<TcpStream>>,
    channels: Arc<SharedChannels>,
) {
    let (write, read) = ws.split();
    let (pwrite, mut pread) = pws.split();

    channels.lock().await.insert(id, write);

    let channels_cloned = channels.clone();
    tokio::spawn(async move {
        while let Some(Ok(msg)) = pread.next().await {
            let msg = map_to_axum_ws_message(msg);
            for write in channels_cloned.lock().await.values_mut() {
                let _ = write.send(msg.clone()).await;
            }
        }
    });

    let _ = read
        .map_ok(map_from_axum_ws_message)
        .map_err(anyhow::Error::from)
        .forward(pwrite.sink_map_err(anyhow::Error::from))
        .await;

    channels.lock().await.remove(&id);
}

async fn ws_student_handler(id: SocketAddr, ws: WebSocket, channels: Arc<SharedChannels>) {
    let (write, mut read) = ws.split();
    channels.lock().await.insert(id, write);
    while let Some(Ok(..)) = read.next().await {}
    channels.lock().await.remove(&id);
}
